from __future__ import annotations

import datetime
from typing import Optional, Dict, Tuple, Type, Union, List, OrderedDict

import numpy as np
from matplotlib import colors
from peewee import FloatField, CharField, Query

from common.typing_utils import numeric
from common.geometry.geom import Geom
from common.logger import logger
from catalog.models.base_stac_item import (
    BaseSTACItem,
    extract_base_info,
    BandDefinition,
)


def create_legend():
    legend = {
        "0": "nodata",
        "1": "saturated-or-defective",
        "2": "dark-area-pixels",
        "3": "cloud-shadows",
        "4": "vegtation",
        "5": "not-vegetated",
        "6": "water",
        "7": "unclassified",
        "8": "cloud-medium-prob",
        "9": "cloud-high-prob",
        "10": "thin-cirrus",
        "11": "snow",
    }

    lst_colors = np.array(
        [
            (0, 0, 0, 0),  # 0: nodata
            (255, 0, 14, 1),  # 1: saturated-or-defective
            (56, 56, 56, 1),  # 2: dark-area-pixels
            (122, 50, 20, 1),  # 3: cloud-shadows
            (0, 255, 60, 1),  # 4: vegtation
            (255, 254, 62, 1),  # 5: not-vegetated
            (0, 16, 191, 1),  # 6: water
            (106, 102, 102, 1),  # 7: unclassified
            (165, 160, 160, 1),  # 8: cloud-medium-prob
            (202, 199, 199, 1),  # 9: cloud-high-prob
            (0, 199, 251, 1),  # 10: thin-cirrus
            (255, 88, 248, 1),  # 11: snow
        ],
        dtype=np.float32,
    )

    lst_colors[:, :-1] = lst_colors[:, :-1] / 255.0
    cmap_scl = colors.ListedColormap(lst_colors, N=lst_colors.shape[0])
    return {"legend": legend, "cmap": cmap_scl}


# legends for the Sentinel-2 L2A SCL data
KNOWN_LEGENDS = {"SCL": create_legend()}


class Sentinel2L2AScene(BaseSTACItem):
    # NOTE to mimic Google Earth Engine's way of defining a product's name
    collection_id = "COPERNICUS/SENTINEL2/L2A"

    bands = OrderedDict()
    bands["B01"] = BandDefinition(
        "B01", "B01", 0, 60, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B02"] = BandDefinition(
        "B02", "B02", 0, 10, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B03"] = BandDefinition(
        "B03", "B03", 0, 10, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B04"] = BandDefinition(
        "B04", "B04", 0, 10, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B05"] = BandDefinition(
        "B05", "B05", 0, 20, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B06"] = BandDefinition(
        "B06", "B06", 0, 20, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B07"] = BandDefinition(
        "B07", "B07", 0, 20, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B08"] = BandDefinition(
        "B08", "B08", 0, 10, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B8A"] = BandDefinition(
        "B8A", "B8A", 0, 20, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B09"] = BandDefinition(
        "B09", "B09", 0, 60, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B11"] = BandDefinition(
        "B11", "B11", 0, 20, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["B12"] = BandDefinition(
        "B12", "B12", 0, 20, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["AOT"] = BandDefinition(
        "AOT", "AOT", 0, 10, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["WVP"] = BandDefinition(
        "WVP", "WVP", 0, 10, 0, 0.0001, 0, np.uint16, (0, 4000)
    )
    bands["SCL"] = BandDefinition("SCL", "SCL", 0, 20, 0, 1, 0, np.uint8, None)

    supported_vis = {
        "rgb": ["B04", "B03", "B02"],
        "nrg": ["B08", "B04", "B03"],
        "legend-scl": ("SCL",),
    }
    default_resolutions = {"meter": 10, "degree": 0.000089831528412}

    # fields
    mgrs_tile = CharField(max_length=5, index=True)
    cloud_cover = FloatField(index=True)

    @classmethod
    def create_one_item(cls: Type[Sentinel2L2AScene], info_item: Dict):
        base_info = extract_base_info(info_item)
        cloud_cover = info_item["properties"]["eo:cloud_cover"]
        mgrs_tile = info_item["properties"]["paii-misc:mgrs_tile"]
        return Sentinel2L2AScene(
            mgrs_tile=mgrs_tile, cloud_cover=cloud_cover, **base_info
        )

    @classmethod
    def construct_query(
        cls: Type[Sentinel2L2AScene],
        spatial: Optional[Geom] = None,
        temporal: Optional[
            Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]
        ] = None,
        created: Optional[
            Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]
        ] = None,
        include_invalidated: bool = False,
        mgrs_tile: Optional[Union[str, List[str]]] = None,
        cloud_cover: Optional[Tuple[Optional[numeric], Optional[numeric]]] = None,
    ) -> Tuple[Query, int]:
        """Query to get many items from the database.

        Parameters
        ----------
        spatial : Optional[Geom], optional
            Area of interest represented by a Geom object, by default None not to filter by spatial area.
        temporal : Optional[Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]], optional
            Temporal range of item (capture time), as a tuple of start and end datetime, or datetime string,
            by default None not to filter by temporal range.
        created : Optional[Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]], optional
            Temporal range of item (created period), as a tuple of start and end datetime, or datetime string,
            by default None not to filter by temporal range.
        include_invalidated : bool, optional
            Whether to include invalidated items in the query result, by default False not to include.
        mgrs_tile : Optional[int], optional
            Maximum number of items to be returned by the query, by default None
        cloud_cover: Optional[str], optional

        Yields
        ------
        Generator[BaseItem, None, None]
            A generator of BaseItem objects that match the query.
        """

        query, n_conditions = super(Sentinel2L2AScene, cls).construct_query(
            spatial=spatial, temporal=temporal, include_invalidated=include_invalidated
        )

        # NOTE only fields with defined index can be queried here
        # for other fields, recommend to use client side post filtering
        if mgrs_tile is not None:
            if isinstance(mgrs_tile, str):
                query = query.where(cls.mgrs_tile == mgrs_tile)
            elif isinstance(mgrs_tile, list):
                query = query.where(cls.mgrs_tile.in_(mgrs_tile))
            n_conditions += 1

        if cloud_cover is not None:
            cloud_cover_lower, cloud_cover_upper = cloud_cover
            if cloud_cover_lower is not None:
                query = query.where(cls.cloud_cover >= cloud_cover_lower)
                n_conditions += 1
            if cloud_cover_upper is not None:
                query = query.where(cls.cloud_cover <= cloud_cover_upper)
                n_conditions += 1

        # temporal queries for created time
        if created is not None:
            start_created, end_created = created
            if start_created is not None:
                query = query.where(cls.created >= start_created)
                n_conditions += 1
            if end_created is not None:
                query = query.where(cls.created <= end_created)
                n_conditions += 1

        return query, n_conditions

    @classmethod
    @BaseSTACItem.auto_connection_value
    def query_many_items(
        cls: Type[Sentinel2L2AScene],
        spatial: Optional[Geom] = None,
        temporal: Optional[
            Tuple[
                Optional[Union[datetime.datetime, str]],
                Optional[Union[datetime.datetime, str]],
            ]
        ] = None,
        created: Optional[
            Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]
        ] = None,
        include_invalidated: bool = False,
        mgrs_tile: Optional[Union[str, List[str]]] = None,
        cloud_cover: Optional[Tuple[Optional[numeric], Optional[numeric]]] = None,
        limit: Optional[int] = None,
        sort_field: Optional[str] = None,
        sort_desc: Optional[bool] = None,
    ) -> List[Optional[Sentinel2L2AScene]]:
        """Query to get many items from the database.

        Parameters
        ----------
        spatial : Optional[Geom], optional
            Area of interest represented by a Geom object, by default None not to filter by spatial area.
        temporal : Optional[Tuple[Optional[datetime.datetime], Optional[datetime.datetime]] ], optional
            Temporal range of item (capture time), as a tuple of start and end datetime, or datetime string,
            by default None not to filter by temporal range.
        created : Optional[Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]], optional
            Temporal range of item (created period), as a tuple of start and end datetime, or datetime string,
            by default None not to filter by temporal range.
        include_invalidated : bool, optional
            Whether to include invalidated items in the query result, by
            default False not to include.
        mgrs_tile : Optional[int], optional
            Maximum number of items to be returned by the query, by default None
        cloud_cover: Optional[str], optional
        limit : Optional[int], optional
            Maximum number of items to be returned by the query, by default None
        sort_field: Optional[str], optional
            Table field used for sorting queries.
        sort_desc: Optional[bool], optional
            Whether to sort in descending order, otherwise ascending.

        Returns
        ------
        List[Optional[BaseStacItem]]
            A list of Sentinel2L2AScene or None.
        """
        query, n_conditions = cls.construct_query(
            spatial=spatial,
            temporal=temporal,
            created=created,
            include_invalidated=include_invalidated,
            mgrs_tile=mgrs_tile,
            cloud_cover=cloud_cover,
        )
        if n_conditions == 0:
            logger.warning("no query condition applied, going to list the whole table")
        if limit is not None:
            query = query.limit(limit)
        if sort_field is not None:
            if sort_desc:
                query = query.order_by(cls._meta.fields[sort_field].desc())
            else:
                query = query.order_by(cls._meta.fields[sort_field])

        items = [item for item in query]
        return items


if __name__ == "__main__":
    import fire

    fire.Fire(
        {
            "make-table": Sentinel2L2AScene.make_table,
            "remove-table": Sentinel2L2AScene.remove_table,
        }
    )
